{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "94a43d53-15b8-4409-972c-6359515bdc97",
   "metadata": {},
   "source": [
    "# 第一步搭建分割模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "85bde612-d786-4274-a69a-98b90eaa0864",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from torchvision import models\n",
    "\n",
    "def convrelu(in_channels, out_channels, kernel, padding):\n",
    "    return nn.Sequential(\n",
    "        nn.Conv2d(in_channels, out_channels, kernel, padding = padding),\n",
    "        nn.ReLU(inplace = True),\n",
    "    )\n",
    "\n",
    "class ResNetUNet(nn.Module):\n",
    "    def __init__(self, n_class):\n",
    "        super().__init__()\n",
    "\n",
    "        self.base_model = models.resnet18(pretrained = True)\n",
    "        self.base_layers = list(self.base_model.children())\n",
    "        \n",
    "        self.layer0 = nn.Sequential(*self.base_layers[:3])\n",
    "        self.layer0_1x1 = convrelu(64, 64, 1, 0)\n",
    "        self.layer1 = nn.Sequential(*self.base_layers[3:5])\n",
    "        self.layer1_1x1 = convrelu(64, 64, 1, 0)\n",
    "        self.layer2 = self.base_layers[5]\n",
    "        self.layer2_1x1 = convrelu(128, 128, 1, 0)\n",
    "        self.layer3 = self.base_layers[6]\n",
    "        self.layer3_1x1 = convrelu(256, 256, 1, 0)\n",
    "        self.layer4 = self.base_layers[7]\n",
    "        self.layer4_1x1 = convrelu(512, 512, 1, 0)\n",
    "\n",
    "        self.upsample = nn.Upsample(scale_factor = 2, mode = 'bilinear', align_corners = True)\n",
    "        self.conv_up3 = convrelu(256 + 512, 512, 3, 1)\n",
    "        self.conv_up2 = convrelu(128 + 512, 256, 3, 1)\n",
    "        self.conv_up1 = convrelu(64 + 256, 256, 3, 1)\n",
    "        self.conv_up0 = convrelu(64 + 256, 128, 3, 1)\n",
    "        self.conv_original_size0 = convrelu(3, 64, 3, 1)\n",
    "        self.conv_original_size1 = convrelu(64, 64, 3, 1)\n",
    "        self.conv_original_size2 = convrelu(64 + 128, 64, 3, 1)\n",
    "        self.conv_last = nn.Conv2d(64, n_class, 1)\n",
    "\n",
    "    def forward(self, input):\n",
    "        x_original = self.conv_original_size0(input)\n",
    "        x_original = self.conv_original_size1(x_original)\n",
    "\n",
    "        layer0 = self.layer0(input)\n",
    "        layer1 = self.layer1(layer0)\n",
    "        layer2 = self.layer2(layer1)\n",
    "        layer3 = self.layer3(layer2)\n",
    "        layer4 = self.layer4(layer3)\n",
    "\n",
    "        layer4 = self.layer4_1x1(layer4)\n",
    "        x = self.upsample(layer4)\n",
    "        layer3 = self.layer3_1x1(layer3)\n",
    "        x = torch.cat([x, layer3], dim = 1)\n",
    "        x = self.conv_up3(x)\n",
    "\n",
    "        x = self.upsample(x)\n",
    "        layer2 = self.layer2_1x1(layer2)\n",
    "        x = torch.cat([x, layer2], dim = 1)\n",
    "        x = self.conv_up2(x)\n",
    "\n",
    "        x = self.upsample(x)\n",
    "        layer1 = self.layer1_1x1(layer1)\n",
    "        x = torch.cat([x, layer1], dim = 1)\n",
    "        x = self.conv_up1(x)\n",
    "\n",
    "        x = self.upsample(x)\n",
    "        layer0 = self.layer0_1x1(layer0)\n",
    "        x = torch.cat([x, layer0], dim = 1)\n",
    "        x = self.conv_up0(x)\n",
    "\n",
    "        x = self.upsample(x)\n",
    "        x = torch.cat([x, x_original], dim = 1)\n",
    "        x = self.conv_original_size2(x)\n",
    "\n",
    "        out = self.conv_last(x)\n",
    "\n",
    "        return out"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "274963a7-e816-4cb2-a6e6-4dfa104c67bc",
   "metadata": {},
   "source": [
    "# 第二步数据加载预处理"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "b9fd2df6-217f-45aa-8391-696027390a60",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from PIL import Image\n",
    "import torchvision.transforms as transforms\n",
    "\n",
    "transform = transforms.Compose([\n",
    "    transforms.Resize((256, 256)),\n",
    "    transforms.ToTensor()\n",
    "])\n",
    "\n",
    "class MyDataset(Dataset):\n",
    "    def __init__(self, img_dir, mask_dir, transform = None):\n",
    "        self.img_dir = img_dir\n",
    "        self.mask_dir = mask_dir\n",
    "        self.transform = transform\n",
    "        self.image_filenames = [f for f in os.listdir(img_dir) if f.endswith('.jpg')]\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.image_filenames)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        img_name = self.image_filenames[idx]\n",
    "        img_path = os.path.join(self.img_dir, img_name)\n",
    "        mask_path = os.path.join(self.mask_dir, img_name)\n",
    "\n",
    "        image = Image.open(img_path).convert(\"RGB\")\n",
    "        mask = Image.open(mask_path).convert(\"L\")\n",
    "\n",
    "        if self.transform:\n",
    "            image = self.transform(image)\n",
    "            mask = self.transform(mask)\n",
    "\n",
    "            mask = (mask > 0).float()\n",
    "        return image, mask\n",
    "\n",
    "train_dataset = MyDataset(img_dir = r'D:\\Car_segmentation\\data\\imgs\\train', mask_dir = r\"D:\\Car_segmentation\\data\\masks\\train\", transform = transform)\n",
    "val_dataset = MyDataset(img_dir = r'D:\\Car_segmentation\\data\\imgs\\val', mask_dir = r\"D:\\Car_segmentation\\data\\masks\\val\", transform = transform)\n",
    "\n",
    "train_loader = DataLoader(train_dataset, batch_size = 8, shuffle = True)\n",
    "val_loader = DataLoader(val_dataset, batch_size = 8, shuffle = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b1759ee5-57ea-40c6-baba-b7b91862a73a",
   "metadata": {},
   "source": [
    "# 第三步，训练参数设置"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "7dafb592-aa9f-438b-8964-0f4fb5997c8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.optim as optim\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "model = ResNetUNet(n_class = 1).to(device)\n",
    "criterion = nn.BCEWithLogitsLoss()\n",
    "optimizer = optim.Adam(model.parameters(), lr = 1e-4)\n",
    "\n",
    "num_epochs = 10"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db438697-69a8-4c6f-b7c4-7d7472495c8a",
   "metadata": {},
   "source": [
    "# 第四步，模型训练"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c735f63b-276d-4fef-b4a6-4acd3eb9d718",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "test: 1\n"
     ]
    }
   ],
   "source": [
    "for epoch in range(num_epochs):\n",
    "    print('test:', epoch + 1)\n",
    "    model.train()\n",
    "    running_loss = 0.0\n",
    "    for index,(inputs, masks) in enumerate(train_loader):\n",
    "        inputs, masks = inputs. to(device), masks.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        outputs = model(inputs)\n",
    "        loss = criterion(outputs, masks)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        running_loss += loss.item()\n",
    "    print(f\"Epoch[{epoch + 1}/{num_epochs}], Loss:{running_loss / len(train_loader):.4f}\")\n",
    "    model.eval()\n",
    "    val_loss=0.0\n",
    "    with torch.no_grad():\n",
    "        for inputs, masks in val_loader:\n",
    "            inputs,masks = inputs.to(device), masks.to(device)\n",
    "            outputs = model(inputs)\n",
    "            loss = criterion(outputs, masks)\n",
    "            val_loss += loss.item()\n",
    "\n",
    "        print(f\"Validation Loss:{val_loss/ len(val_loader):.4f}\")\n",
    "torch.save(model.state_dict(),'Last_model.pth')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51ac1360-e150-4bb3-b854-b70192cb6d89",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af2c4788-72be-4d36-9faf-74ccd6a2d78b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35c60ab8-c7fe-41b3-94fa-5e7505a83c4c",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
